[Perf] Add new npu_fused_infer_attention_score op to improve perfomance in splitfuse cases and resolve long-seq mask problems #2962

qyqc731 · 2025-09-16T11:56:46Z

What this PR does / why we need it?

Add new npu_fused_infer_attention_score op to improve perfomance in splitfuse cases and resolve long-seq mask problems .

The original op's performance is suboptimal in certain scenarios, necessitating optimization through the new op (npu_fused_infer_attention_score)。
For ultra-long sequences (128k), the original operator will allocate a large attn_mask, which consumes excessive CPU memory. In contrast, the new op supports a fixed-size compressed mask, effectively resolving this issue.

NOTE1: The current PR retains the original logic and uses a version check of the CANN package to determine whether the new op can be enabled. This ensures no impact on existing users. In future versions, this version check and the original logic will be deprecated, and the new op scheduling will be uniformly adopted.
NOTE2: This pr relies on future CANN version, which is not available now.
NOTE3: To enable the new op in chunked prefill, the parameter additional_config should be set like --additional-config '{"ascend_scheduler_config": {"enabled":true,"enable_chunked_prefill":true}}' \ at least.

Does this PR introduce any user-facing change?

No

How was this patch tested?

CI passed

vLLM version: v0.10.2
vLLM main: vllm-project/vllm@6c5f82e

github-actions · 2025-09-16T11:56:57Z

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

A PR should do only one thing, smaller PRs enable faster reviews.
Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

gemini-code-assist

Code Review

这个 PR 主要是为了接入新的 splitfuse chunked prefill 算子。代码改动涉及 attention_v1.py 和 model_runner_v1.py 两个文件。在 attention_v1.py 中，_forward_v1_style 函数的注意力计算从 _npu_paged_attention_splitfuse 切换到了 npu_fused_infer_attention_score。在 model_runner_v1.py 中，为 ChunkedPrefill 场景生成 attention mask 的逻辑被修改。

我的审查发现两个严重问题：

在 attention_v1.py 中，传递给新算子的 actual_seq_lengths 参数值是错误的，使用了累积的 token 位置而不是序列长度，这会导致注意力计算错误。
在 model_runner_v1.py 中，为 ChunkedPrefill 生成的 attention mask 使用了硬编码的尺寸 (2048, 2048)，这使得代码很脆弱，当序列长度超过 2048 时会导致错误。

建议修复这两个严重问题以保证代码的正确性和健壮性。

gemini-code-assist · 2025-09-16T11:59:11Z

vllm_ascend/attention/attention_v1.py

-            num_kv_heads=self.num_kv_heads,
+            input_layout="TND",
+            block_size=block_size,
+            actual_seq_lengths=attn_metadata.query_start_loc[1:],


actual_seq_lengths 参数被传递了 attn_metadata.query_start_loc[1:]，它包含的是累积的 token 位置，而不是各个序列的长度。这很可能会导致错误的注意力计算结果。你应该使用 attn_metadata.query_lens，它才包含正确的序列长度，并确保它在正确的设备上。

Suggested change

actual_seq_lengths=attn_metadata.query_start_loc[1:],

actual_seq_lengths=attn_metadata.query_lens.to(query.device),

plz consider use attn_metadata.query_lens

gemini-code-assist · 2025-09-16T11:59:11Z

vllm_ascend/worker/model_runner_v1.py

        if attn_state == AscendAttentionState.ChunkedPrefill and not self.vllm_config.model_config.use_mla:
-            return self.attn_mask_builder.get_splitfuse_attn_mask(
-                seq_lens, position, self.dtype, self.device)
+            return torch.triu(torch.ones(2048, 2048), diagonal=1).to(torch.int8)


Attention mask 使用了硬编码的尺寸 (2048, 2048)。这是一个魔法数字，使得实现不够健壮。如果批处理中任何序列的长度超过 2048，将导致不正确的掩码或越界错误。掩码的大小应该由模型配置的最大序列长度决定，以确保正确性并避免魔法数字。

Suggested change

return torch.triu(torch.ones(2048, 2048), diagonal=1).to(torch.int8)

return torch.triu(torch.ones(self.model_config.max_model_len, self.model_config.max_model_len), diagonal=1).to(torch.int8)

this is because new op support compressed mask

rjg-lyh · 2025-09-17T07:36:05Z

Could you integrate other scenarios, such as full FlashAttention, using the FIA interface as well, and provide the performance test results?

Angazenn · 2025-09-19T02:40:39Z

vllm_ascend/worker/model_runner_v1.py


        self.attn_mask_builder = AttentionMaskBuilder(
-            self.model_config.max_model_len, self.dtype)
+            self.model_config.max_model_len, self.dtype, self.device)


If we pass self.model_config.max_model_len here we still have an attention mask of shape [max_model_len, max_model_len]. We may have 2 choices:

Passing self.scheduler_config.max_num_batched_tokens instead of self.model_config.max_model_len.

Also using new attention op for full prefill case.

Angazenn · 2025-09-19T02:41:25Z

vllm_ascend/attention/attention_v1.py

-            num_kv_heads=self.num_kv_heads,
+            input_layout="TND",
+            block_size=block_size,
+            actual_seq_lengths=attn_metadata.query_start_loc[1:],


plz consider use attn_metadata.query_lens

Angazenn · 2025-09-19T03:04:21Z

Also notice another 2 points:

compatibility with difference torch_npu/cann version.
fix bug when padding is introduced in dp (len of query exceeds attn_metadata.query_lens[-1]).

github-actions · 2025-09-19T07:51:03Z

This pull request has conflicts, please resolve those before we can evaluate the pull request.

qyqc731 · 2025-09-19T08:13:30Z

vllm_ascend/attention/attention_v1.py

-            num_heads=self.num_heads,
-            scale_value=self.scale,
-            out=output)
+        if self.compressed_mask:


add TODO: This op will be used in more situation in future.

weijinqian0 · 2025-09-19T08:15:14Z

vllm_ascend/worker/model_runner_v1.py

-            return self.attn_mask_builder.get_splitfuse_attn_mask(
-                seq_lens, position, self.dtype, self.device)
+            if selkf.compressed_mas:
+                return self.attn_mask_builder.get_splitfuse_attn_mask()


qyqc731 · 2025-09-19T08:17:50Z

vllm_ascend/utils.py

    raise ValueError("Invalid type for tensors")
+
+
+def verify_torch_npu_version(


The cann version needs to be verfied. using package.version

add todo for temp

qyqc731 · 2025-09-19T08:20:15Z

vllm_ascend/attention/attention_mask.py

+            # so device needs to be passed here.
+            assigned_mask_dim = 2048
+            self.chunked_prefill_attn_mask = torch.triu(torch.ones(assigned_mask_dim, assigned_mask_dim), diagonal=1
+            ).to(torch.int8).to(device)


check type again

this op currently only supports int8 mask

vllm_ascend/attention/attention_mask.py

Angazenn · 2025-09-19T02:59:47Z

vllm_ascend/worker/model_runner_v1.py


        self.attn_mask_builder = AttentionMaskBuilder(
-            self.model_config.max_model_len, self.dtype)
+            self.model_config.max_model_len, self.dtype, self.device)


for choice 1:

Suggested change

self.model_config.max_model_len, self.dtype, self.device)

self.scheduler_config.max_num_batched_tokens, self.dtype, self.device)

Angazenn · 2025-09-19T03:07:26Z

vllm_ascend/worker/model_runner_v1.py

        if attn_state == AscendAttentionState.ChunkedPrefill and not self.vllm_config.model_config.use_mla:
-            return self.attn_mask_builder.get_splitfuse_attn_mask(
-                seq_lens, position, self.dtype, self.device)
+            return torch.triu(torch.ones(2048, 2048), diagonal=1).to(torch.int8)


this is because new op support compressed mask

Angazenn · 2025-09-19T14:40:40Z

vllm_ascend/attention/attention_mask.py

+            # so device needs to be passed here.
+            assigned_mask_dim = 2048
+            self.chunked_prefill_attn_mask = torch.triu(torch.ones(assigned_mask_dim, assigned_mask_dim), diagonal=1
+            ).to(torch.int8).to(device)


this op currently only supports int8 mask

Signed-off-by: tangtianyi <tangtianyi4@huawei.com>

Signed-off-by: Angazenn <supperccell@163.com>

MengqingCao

Thanks for this work! just note: we need to remember to upgrade torch-npu and cann together in ci later

…ce in splitfuse cases and resolve long-seq mask problems (vllm-project#2962) ### What this PR does / why we need it? Add new npu_fused_infer_attention_score op to improve perfomance in splitfuse cases and resolve long-seq mask problems . 1. The original op's performance is suboptimal in certain scenarios, necessitating optimization through the _new op_ (npu_fused_infer_attention_score)。 2. For ultra-long sequences (128k), the original operator will allocate a large attn_mask, which consumes excessive CPU memory. In contrast, the _new op_ supports a fixed-size compressed mask, effectively resolving this issue. NOTE1: The current PR retains the original logic and uses a version check of the CANN package to determine whether the _new op_ can be enabled. This ensures no impact on existing users. In future versions, this version check and the original logic will be deprecated, and the _new op_ scheduling will be uniformly adopted. NOTE2: This pr relies on future CANN version, which is not available now. NOTE3: To enable the new op in chunked prefill, the parameter additional_config should be set like `--additional-config '{"ascend_scheduler_config": {"enabled":true,"enable_chunked_prefill":true}}' \` at least. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? CI passed - vLLM version: v0.10.2 - vLLM main: vllm-project/vllm@6c5f82e --------- Signed-off-by: tangtianyi <tangtianyi4@huawei.com> Signed-off-by: Angazenn <supperccell@163.com> Co-authored-by: Angazenn <supperccell@163.com> Signed-off-by: Che Ruan <cr623@ic.ac.uk>

qyqc731 changed the title ~~chunked prefill splitfuse算子接入~~ chunked prefill splitfuse op in Sep 16, 2025

gemini-code-assist bot reviewed Sep 16, 2025

View reviewed changes

qyqc731 changed the title ~~chunked prefill splitfuse op in~~ chunked prefill, access splitfuse op Sep 16, 2025

Angazenn reviewed Sep 19, 2025

View reviewed changes

github-actions bot added the module:core label Sep 19, 2025

github-actions bot added the merge-conflicts label Sep 19, 2025

qyqc731 commented Sep 19, 2025

View reviewed changes

weijinqian0 reviewed Sep 19, 2025

View reviewed changes

qyqc731 commented Sep 19, 2025

View reviewed changes

MengqingCao reviewed Sep 19, 2025

View reviewed changes

vllm_ascend/attention/attention_mask.py Show resolved Hide resolved

Angazenn force-pushed the main branch from b8515ab to e2e335b Compare September 19, 2025 14:33

Angazenn reviewed Sep 19, 2025

View reviewed changes

github-actions bot removed the merge-conflicts label Sep 19, 2025

qyqc731 and others added 8 commits September 19, 2025 22:56

chunked prefill splitfuse算子接入

66923d7

Signed-off-by: tangtianyi <tangtianyi4@huawei.com>

splitfuse access optimize

b0eaf77

Signed-off-by: tangtianyi <tangtianyi4@huawei.com>

optimize again

613c927

Signed-off-by: tangtianyi <tangtianyi4@huawei.com>

improve compatibility

a678de9

Signed-off-by: tangtianyi <tangtianyi4@huawei.com>

fix for dp padding

7020e3a

Signed-off-by: tangtianyi <tangtianyi4@huawei.com>

add compatibility for attention v1

8bdda1d

Signed-off-by: tangtianyi <tangtianyi4@huawei.com>

fix by review

6ef548d

Signed-off-by: tangtianyi <tangtianyi4@huawei.com>

change to cann version

821253c

Signed-off-by: Angazenn <supperccell@163.com>

Angazenn force-pushed the main branch from 026b1a5 to 821253c Compare September 19, 2025 14:57

Angazenn added 2 commits September 19, 2025 23:06

remove another compressed

35acbb2

Signed-off-by: Angazenn <supperccell@163.com>

fix lint

a366af6

Signed-off-by: Angazenn <supperccell@163.com>

github-actions bot removed the module:core label Sep 20, 2025

Angazenn added 4 commits September 20, 2025 10:23

fix mypy

5102efe

Signed-off-by: Angazenn <supperccell@163.com>

fix mypy

5565a33

Signed-off-by: Angazenn <supperccell@163.com>

fix mypy

8ac3f51

Signed-off-by: Angazenn <supperccell@163.com>

fix ut

3be8a40

Signed-off-by: Angazenn <supperccell@163.com>

MengqingCao added ready read for review ready-for-test start test by label for PR labels Sep 20, 2025

wangxiyuan added ready-for-test start test by label for PR and removed ready-for-test start test by label for PR labels Sep 22, 2025

qyqc731 changed the title ~~chunked prefill, access splitfuse op~~ [Perf] Add new npu_fused_infer_attention_score op to improve perfomance in splitfuse cases and resolve long-seq mask problems Sep 22, 2025

MengqingCao approved these changes Sep 22, 2025

View reviewed changes

MengqingCao merged commit f1f2c8f into vllm-project:main Sep 22, 2025
59 of 62 checks passed

	actual_seq_lengths=attn_metadata.query_start_loc[1:],
	actual_seq_lengths=attn_metadata.query_lens.to(query.device),

	return torch.triu(torch.ones(2048, 2048), diagonal=1).to(torch.int8)
	return torch.triu(torch.ones(self.model_config.max_model_len, self.model_config.max_model_len), diagonal=1).to(torch.int8)

		raise ValueError("Invalid type for tensors")


		def verify_torch_npu_version(

	self.model_config.max_model_len, self.dtype, self.device)
	self.scheduler_config.max_num_batched_tokens, self.dtype, self.device)

[Perf] Add new npu_fused_infer_attention_score op to improve perfomance in splitfuse cases and resolve long-seq mask problems #2962

[Perf] Add new npu_fused_infer_attention_score op to improve perfomance in splitfuse cases and resolve long-seq mask problems #2962

Conversation

qyqc731 commented Sep 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What this PR does / why we need it?

Does this PR introduce any user-facing change?

How was this patch tested?

Uh oh!

github-actions bot commented Sep 16, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Sep 16, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Sep 16, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rjg-lyh commented Sep 17, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Angazenn commented Sep 19, 2025

Uh oh!

github-actions bot commented Sep 19, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MengqingCao left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

qyqc731 commented Sep 16, 2025 •

edited

Loading